#ifndef __MODULE_MODEL_LOADER_H_
#define __MODULE_MODEL_LOADER_H_

#include <iostream>
#include <string>
#include <vector>

#include "standard_api.h"
#include "common/dbg.hpp"

class ModelLoader {
 public:
  ModelLoader(std::string filename) : filename_(filename) {
    std::cout << filename << std::endl;
  }

  ~ModelLoader() {
    if (ctx_) {
      int ret;
      if (graph_id_) {
        ret = aipu_unload_graph(ctx_, graph_id_);
        if (ret != AIPU_STATUS_SUCCESS) {
          const char* msg = nullptr;
          aipu_get_error_message(ctx_, static_cast<aipu_status_t>(ret), &msg);
          AIPU_ERR()("aipu_unload_graph: %s\n", msg);
        } else {
          AIPU_INFO()("aipu_unload_graph success\n");
        }
      }

      ret = aipu_deinit_context(ctx_);
      if (ret != AIPU_STATUS_SUCCESS) {
        const char* msg = nullptr;
        aipu_get_error_message(ctx_, static_cast<aipu_status_t>(ret), &msg);
        AIPU_ERR()("aipu_deinit_ctx: %s\n", msg);
      }
    }
  }

  int Init() {
    if (inited_ == true) return 0;
    aipu_status_t ret = aipu_init_context(&ctx_);
    const char* msg = nullptr;
    if (ret != AIPU_STATUS_SUCCESS) {
      aipu_get_error_message(ctx_, ret, &msg);
      AIPU_ERR()("aipu_init_context: %s\n", msg);
      return -1;
    }

    ret = aipu_load_graph(ctx_, filename_.c_str(), &graph_id_);
    std::cout << "graph_id_:" << graph_id_ << std::endl;
    if (ret != AIPU_STATUS_SUCCESS) {
      aipu_get_error_message(ctx_, ret, &msg);
      AIPU_ERR()("aipu_load_graph_helper: %s (%s)\n", msg, filename_);
      return -1;
    }

    ret = aipu_get_tensor_count(ctx_, graph_id_, AIPU_TENSOR_TYPE_INPUT, &input_cnt_);
    if (ret != AIPU_STATUS_SUCCESS) {
      aipu_get_error_message(ctx_, ret, &msg);
      AIPU_ERR()("aipu_get_tensor_count: %s\n", msg);
      return -1;
    }
    for (uint32_t i = 0; i < input_cnt_; i++) {
      aipu_tensor_desc_t desc;
      ret = aipu_get_tensor_descriptor(ctx_, graph_id_, AIPU_TENSOR_TYPE_INPUT, i, &desc);
      if (ret != AIPU_STATUS_SUCCESS) {
        aipu_get_error_message(ctx_, ret, &msg);
        AIPU_ERR()("aipu_get_tensor_descriptor: %s\n", msg);
        return -1;
      }
      input_desc_.push_back(desc);
    }

    ret = aipu_get_tensor_count(ctx_, graph_id_, AIPU_TENSOR_TYPE_OUTPUT, &output_cnt_);
    if (ret != AIPU_STATUS_SUCCESS) {
      aipu_get_error_message(ctx_, ret, &msg);
      AIPU_ERR()("aipu_get_tensor_count: %s\n", msg);
      return -1;
    }
    for (uint32_t i = 0; i < output_cnt_; i++) {
      aipu_tensor_desc_t desc;
      ret = aipu_get_tensor_descriptor(ctx_, graph_id_, AIPU_TENSOR_TYPE_OUTPUT, i, &desc);
      if (ret != AIPU_STATUS_SUCCESS) {
        aipu_get_error_message(ctx_, ret, &msg);
        AIPU_ERR()("aipu_get_tensor_descriptor: %s\n", msg);
      }
      output_desc_.push_back(desc);
    }
    inited_ = true;
    return 0;
  }

  aipu_ctx_handle_t* GetContext() {
    return ctx_;
  }
  uint64_t GetGraphId() {
    return graph_id_;
  }
  size_t GetInputCnt() {
    return input_cnt_;
  }
  size_t GetOutputCnt() {
    return output_cnt_;
  }
  std::vector<aipu_tensor_desc_t> GetInputDesc() {
    return input_desc_;
  }
  std::vector<aipu_tensor_desc_t> GetOutputDesc() {
    return output_desc_;
  }

 private:
  uint64_t graph_id_ = 0;
  std::string filename_ = "";
  aipu_ctx_handle_t* ctx_ = nullptr;

 private:
  bool inited_ = false;
  uint32_t input_cnt_;
  uint32_t output_cnt_;
  std::vector<aipu_tensor_desc_t> input_desc_;
  std::vector<aipu_tensor_desc_t> output_desc_;
};

#endif
